
import torch
import numpy as np
from sklearn import metrics
from torch.utils.data import DataLoader
from torchvision import transforms

from .AttackUtil import *
from .DataProcess import *

def get_loss(dataloader, loss_fn, device):
    i = 0
    for batch, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)
        loss = loss_fn(X.to(torch.float64), y.to(torch.int64))
        if i == 0:
            loss_list = loss
        else:
            loss_list = torch.cat((loss_list, loss), 0)
        i += 1
    return loss_list

def loss_attack_evaluate(loss_list, mem_list):
    average_loss = np.mean(loss_list)
    median_loss = np.median(loss_list)
    auc_value = metrics.roc_auc_score(mem_list, -loss_list)
    print("AUC value is:", auc_value)
    accuracy = metrics.accuracy_score(mem_list, (loss_list < median_loss).astype(int))
    print("Accuracy is:", accuracy)


def conf_attack_evaluate(max_conf_list, mem_list):
    average_max_conf = np.mean(max_conf_list)
    median_max_conf = np.median(max_conf_list)
    auc_value = metrics.roc_auc_score(mem_list, max_conf_list)
    print("AUC value is:", auc_value)
    accuracy = metrics.accuracy_score(mem_list, (max_conf_list > average_max_conf).astype(int))
    print("Accuracy is:", accuracy)

# 成员推断攻击——基于损失的阈值攻击
def loss_threshold_attack(tar_model, weight_dir, data_name, model, model_transform, model_epochs, batch_size, loss_fn, device, prop_keep=0.14):
    # 加载目标模型数据
    conf_data_target, m_data_target, label_data_target = load_attack_data(weight_dir, tar_model, data_name, model, model_epochs, model_transform, device, prop_keep=prop_keep)
    print(conf_data_target.shape, m_data_target.shape, label_data_target.shape)

    transform = transforms.Compose([])
    loss_dataset_target = CustomDataset(conf_data_target, label_data_target, transform)
    loss_dataloader_target = DataLoader(loss_dataset_target, batch_size=batch_size, shuffle=False)

    loss_list = get_loss(loss_dataloader_target, loss_fn, device)
    loss_list = loss_list.detach().cpu().numpy()

    loss_attack_evaluate(loss_list, m_data_target)


def conf_threshold_attack(tar_model, weight_dir, data_name, model, model_transform, model_epochs, device, prop_keep=0.14):

    # 加载目标模型数据
    conf_data_target, m_data_target, label_data_target = load_attack_data(weight_dir, tar_model, data_name, model, model_epochs, model_transform, device, prop_keep=prop_keep)
    print(conf_data_target.shape, m_data_target.shape, label_data_target.shape)

    max_conf_list = np.amax(conf_data_target, axis=1)

    conf_attack_evaluate(max_conf_list, m_data_target)


def loss_threshold_attack_simp(conf_data_target, m_data_target, label_data_target, batch_size, loss_fn, device):
    
    transform = transforms.Compose([])
    loss_dataset_target = CustomDataset(conf_data_target, label_data_target, transform)
    loss_dataloader_target = DataLoader(loss_dataset_target, batch_size=batch_size, shuffle=False)

    loss_list = get_loss(loss_dataloader_target, loss_fn, device)
    loss_list = loss_list.detach().cpu().numpy()

    loss_attack_evaluate(loss_list, m_data_target)


def conf_threshold_attack_simp(conf_data_target, m_data_target, label_data_target):

    max_conf_list = np.amax(conf_data_target, axis=1)
    conf_attack_evaluate(max_conf_list, m_data_target)